
import os
import torch
import importlib.util
from torchvision import models

from config_manager.model_config import ModelConfig


class ModelManager(object):
    """
    模型管理器  模型下载  模型加载
    """

    def __init__(self,config: ModelConfig, device="cuda"):

        self.config = config
        self.device = device
        self.parameter_validation()

    def parameter_validation(self):
        """
        对配置的模型信息进行验证
        """
        if self.config.model_type == "builtin":
            # 1️⃣ 内置模型
            if not hasattr(models, self.config.model_name):
                raise ValueError(f"未找到内置模型：{self.config.model_name}")

        elif self.config.model_type == "custom":
            if not os.path.exists(self.config.custom_model_path):
                raise FileNotFoundError(f"模型文件不存在: {self.config.custom_model_path}")

        else:
            raise ValueError(f"未知的模型类型：{self.config.model_type}")


    def get_model(self):
        if self.config.model_type == "builtin":
            # 1️⃣ 内置模型
            model_fn = getattr(models, self.config.model_name)
            model = model_fn(num_classes=self.config.num_classes)
            return model

        elif self.config.model_type == "custom":

            # 从文件路径加载模块
            spec = importlib.util.spec_from_file_location(
                "custom_model", self.config.custom_model_path
            )
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)

            # 获取模型类
            model_class = getattr(module, self.config.custom_model_class, None)
            if model_class is None:
                raise ValueError(
                    f"在 {self.config.custom_model_path} 中未找到类 {self.config.custom_model_class}"
                )

            # 初始化模型实例
            if self.config.model_args is not None:
                model = model_class(**self.config.model_args).to(device=self.device, dtype=torch.float32)
            else:
                model = model_class().to(device=self.device, dtype=torch.float32)
            return model


